from FitnessFunction import Fitness
import numpy as np
from matplotlib import pyplot as plt
from FitnessPlot import FitnessPlot

################################################################
# 果蝇优化算法
################################################################


class FOA():
    """果蝇优化算法"""

    def __init__(self, FitFunction, LB, UB, dim, population_size, max_iter):
        """初始化果蝇算法的参数"""
        self.FitFunction = FitFunction
        self.LB = LB
        self.UB = UB
        self.dim = dim
        self.population_size = population_size
        self.max_iter = max_iter
        self.X = np.zeros((self.population_size, self.dim))
        self.gBest_score = np.inf  # 最优解对应的适应度值
        self.gBest_curve = np.zeros(self.max_iter)  # 每轮的最优解的适应度值
        self.gBest_X = np.zeros(dim)  # 每次迭代得到的最佳位置
        self.population_fitness = np.zeros(self.population_size)

    def Initialize_population(self):
        """initialize population 初始化种群数据"""
        # 初始化种群数据X
        for i in range(self.population_size):
            for j in range(self.dim):
                self.X[i, j] = (self.UB[j] - self.LB[j]) * \
                    np.random.random() + self.LB[j]

    def BorderCheck_UpdateFitness(self):
        """边界检查与更新适应度值"""
        for i in range(self.population_size):
            for j in range(self.dim):
                # 检查边界
                for j in range(self.dim):
                    if (self.X[i, j] > self.UB[j]):
                        # 不符合上边界
                        self.X[i, j] = self.UB[j]

                    if (self.X[i, j] < self.LB[j]):
                        # 不符合下边界
                        self.X[i, j] = self.LB[j]
                # 计算所有物种的fitness
                fitness = self.CalculateFitness(self.X[i, :])  # 获取物种的适应度值
                # 更新最优解的适应度值和物种的最优解的位置信息
                if fitness < self.gBest_score:
                    self.gBest_score = fitness
                    self.gBest_X = self.X[i, :].copy()

    def CalculateFitness(self, X):
        """自定义计算适应度函数,计算单个物种的适应度值"""
        fitness = self.FitFunction(X)
        return fitness

    def Optimize(self):
        t = 0
        while t < self.max_iter:
            for i in range(self.population_size):
                # 每个个体朝着气味浓度最重的地方（适应度最优的位置）进行移动
                self.X[i, :] = self.gBest_X+2 * np.random.rand() - 1
                self.population_fitness[i] = self.CalculateFitness(
                    self.X[i, :])
            # 对种群历史最优位置信息与适应度值进行更新
            if self.population_fitness.min() < self.gBest_score:
                self.gBest_score = self.population_fitness.min()
                self.gBest_X = self.X[self.population_fitness.argmin(
                ), :].copy()
            # 存储当前迭代下的种群历史最优适应度值并输出
            # 更新每轮的最优解的适应度值
            self.gBest_curve[t] = self.gBest_score
            if (t % 10 == 0):
                print('第{}轮训练完成!, 该轮最优解的适应度值为: {}'.format(
                    t, self.gBest_curve[t]))
            t += 1

        return self.gBest_curve, self.gBest_X


if __name__ == '__main__':
    population_size = 30
    dim = 30
    max_iter = dim * 10000
    LB = -10 * np.ones(dim)
    UB = 10 * np.ones(dim)
    Func = Fitness(min=LB, max=UB, dim=dim)
    FitFunction = Func.F10
    Gbest_curve, GbestX = FOA(FitFunction, LB=LB, UB=UB,
                              dim=dim, population_size=population_size, max_iter=max_iter).Optimize()

    # 可视化训练图像
    FitnessPlot.Plot_2D('FOA', Gbest_curve)
